#ifndef SIMPLE_HAT_HPP
#define SIMPLE_HAT_HPP
/*
======================================================================================
simple_hat Abstract Data Type       Copyright (c) 2010 Jonathan "AngleWyrm" Wooldridge

  This version of the hat class has only the functions necessary to make a hat,
  and is implemented as a container adapter, which defaults to vector.
  Any container that provides the following functions should be suitable:
    operator[] access to data
    push_back()
    pop_back()
    size()

  The public functions of this hat:
  * put   -- inserts an object into the container
  * get   -- returns an object, without removing it
  * pull  -- returns an object, and removes that object from the container
  * size  -- returns count of items contained
  * print -- debugging routine to trace operations

http://code.google.com/p/hat-container/source/browse/branches/SimpleHat/simple_hat.hpp
Permission is hereby granted, free of charge, to any person obtaining a copy of this
software and associated documentation files (the "Software"), to deal in the Software
without restriction, including without limitation the rights to use, copy, modify,
merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so,
subject to the following conditions:

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS
OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE
======================================================================================
*/

//////////////////////////////////////////////////////////////////////////////////////
// Includes
#include <vector>
using std::vector;

#include <iostream>
using std::cout;
using std::endl;

#include <cstdlib>
using std::rand;

namespace hat
{

//////////////////////////////////////////////////////////////////////////////////////
// Node class for container
template<class T>
class node_type
{
public:
    T item;
    size_t chance_weight;    // num chances for this item
    size_t family_weight;    // sum of self and children's chance weights

    node_type( const T& input = T(), const size_t chances = 1)
    : item(input), chance_weight(chances), family_weight(chances) { };
};


//////////////////////////////////////////////////////////////////////////////////////
// Simple Hat
//
template<class T, class container_type = vector<node_type<T> > >
class simple_hat
{
public:
    // Main public interface
    void put(T const &item, const size_t chances = 1);
    T&   get();
    T&   pull(T&);

    // secondary utility functions
    size_t size() const { return tree.size(); }
    void   print(size_t index) const;

protected:
    size_t find_index(size_t index, size_t target);
    void   update_weights(size_t index);
    void   erase(size_t index);

    container_type tree;
};

//////////////////////////////////////////////////////////////////////////////////////
// put
template <class T, class container_type >
void simple_hat<T, container_type>::put(const T& input, const size_t chances )
{
    node_type<T> new_node(input, chances );
    tree.push_back( new_node );

    // update family chance weights
    update_weights( tree.size() - 1 );
};

//////////////////////////////////////////////////////////////////////////////////////
// update weights; used by put()
template <class T, class container_type>
void simple_hat<T, container_type>::update_weights( size_t index )
{
    // update the weight of the currently indexed node
    tree[index].family_weight =  tree[index].chance_weight;

    // First, update self with children's weights
    // daughters are found at 2(x+1)
    if( 2*(index+1) < tree.size() )
    {   // add daughter's weight to family
        tree[index].family_weight += tree[ 2*(index+1) ].family_weight;
    } // else no daughter

    // sons are found at 2x+1
    if( 2*index+1 < tree.size() )
    {   // add son
        tree[index].family_weight += tree[ 2*index+1 ].family_weight;
    } // else no son

    // Second, update parents with self
    // check for parents, and update them
    if( index > 0 ) // not the root mother (has parent)
    {
        // calculate distance to parent
        if( index % 2 ) { // is a male node
            update_weights( (index-1)/2 ); // update parent
        }
        else { // is a female node
            update_weights( index/2 - 1 ); // update parent
        }
    } // else this is the root mother (no parents)
};

//////////////////////////////////////////////////////////////////////////////////////
// get()
template <class T, class container_type>
T& simple_hat<T, container_type>::get()
{
    // roll a random number on the sum of all chance weights
    size_t target = rand() % ( tree[0].family_weight );

    // find the matching index in the tree
    size_t index = find_index(0, target);
    return tree[index].item;
};

//////////////////////////////////////////////////////////////////////////////////////
// find index; used by get() & pull()
template <class T, class container_type>
size_t simple_hat<T, container_type>::find_index( size_t index, size_t target )
{
    // divide weight into three groups: Self, Son & Daughter, and test each

    // Testing self
    size_t sum = tree[index].chance_weight;
    if( target < sum )
    {
        return index; // found within self
    } // else outside self, must be in children

    // Testing son (first-born is always a son)
    sum += tree[ 2*index+1 ].family_weight;
    if( target < sum )
    {
        return find_index( 2*index+1, target - tree[index].chance_weight );
    } // else not in son's range

    // Testing daughter
    sum += tree[ 2*(index+1) ].family_weight;
    if( target < sum )
    {
        return find_index( 2*(index+1),
         target - tree[index].chance_weight - tree[ 2*index+1 ].family_weight );
    }

    // possible to reach here if a custom random number generator is used,
    // and it fails to deliver a number in the correct range
    return 0;
};

//////////////////////////////////////////////////////////////////////////////////////
// pull
template <class T, class container_type>
T& simple_hat<T, container_type>::pull( T& buffer )
{
    // get random item
    size_t random_weight = rand() % tree[0].family_weight;
    size_t index = find_index( 0, random_weight );
    buffer = tree[index].item;

    erase( index );
    return buffer;
};

//////////////////////////////////////////////////////////////////////////////////////
// erase; used by pull
template <class T, class container_type>
void simple_hat<T, container_type>::erase( size_t index )
{

    // erase weights from chance weights tree by overwriting with last node
    tree[index] = tree[ tree.size() - 1 ]; // copy last node over this node
    tree.pop_back();                       // delete old node
    if( size()==0 ) return;                    // if we deleted last node, bail
    update_weights(index);                 // update replacement node's parents

    // update replacement node's old parents
    if( tree.size() % 2 ) { // then odd number, old size was even (female node)
        update_weights( (tree.size()-1)/2 ); // update old parent of female
    }
    else { // was a male node
        update_weights( tree.size()/2 - 1 ); // update old parent of male
    }
};

//////////////////////////////////////////////////////////////////////////////////////
// print, for debugging.
// Only works for data that can be passed to cout
template <class T, class container_type>
void simple_hat<T, container_type>::print(size_t index) const
{
    cout << endl << tree[index].item
         << "("  <<  tree[index].chance_weight
         << "/"  <<  tree[index].family_weight << ")";

    if(2*index+1 < size()){ // has son
        cout << "\tson[" << 2*index+1 << "]:" << tree[2*index+1].item;
    }
    if(2*(index+1) < size()) { // also has daughter
        cout << "\tdaughter[" << 2*(index+1) << "]:" << tree[2*(index+1)].item;
    }
};

}; // namespace hat
#endif // SIMPLE_HAT_HPP defined
